[All] Refactor nvte_get_fused_attn_backend with cudnn-frontend calls by cyanguwa · Pull Request #2964 · NVIDIA/TransformerEngine

cyanguwa · 2026-05-06T01:56:59Z

Description

This PR replaces the hand-maintained backend selection logic in nvte_get_fused_attn_backend with cudnn-frontend's production-grade support checks.

It will build the same graph that runtime uses in execution, cache the graph if the build is successful, and provide a warmed-up cache for the next nvte_get_fused_attn_backend call.
It provides a cleaner dispatch logic, avoids accidental regressions, and helps TE to stay in sync with cudnn-frontend's support surface.
It also provides appropriate error messaging for when no backend is available, so users are notified and make config/architecture/cudnn version adjustments.

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

See Description.

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

for more information, see https://pre-commit.ci

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

for more information, see https://pre-commit.ci

greptile-apps · 2026-05-08T00:17:25Z

Greptile Summary

This PR replaces the hand-maintained boolean logic in nvte_get_fused_attn_backend with cudnn-frontend probe calls: it actually builds the same cuDNN graph that the runtime kernel will use (with null device pointers), caches it for future calls, and reports the real cuDNN error message when no backend is available.

A new FusedAttnConfig/FusedAttnFwdParams/FusedAttnBwdParams opaque-handle API is introduced in the public header, with v2 entry-points (nvte_get_fused_attn_backend_v2, nvte_fused_attn_fwd_v2, nvte_fused_attn_bwd_v2); PyTorch and JAX bindings are updated to use them.
The old nvte_get_fused_attn_backend is kept as a deprecated thin wrapper that constructs a FusedAttnConfig and delegates to the v2 path, but omits batch_size (stays 0), causing cuDNN to reject every probe graph and silently return NVTE_No_Backend for all BSHD/SBHD F16 configs.
graph_debug.h is included unconditionally and registers an atexit handler and atomic counters in every build; it is self-labeled "TEMPORARY DEBUG INSTRUMENTATION — REMOVE AFTER VERIFICATION".

Confidence Score: 2/5

Not safe to merge as-is: the deprecated wrapper regression silently returns NVTE_No_Backend for all BSHD/SBHD F16 configs, and temporary debug instrumentation with an atexit handler ships unconditionally.

The probe-based v2 approach is architecturally sound, but two defects block a clean merge: (1) the deprecated wrapper omits batch_size, causing cuDNN to reject every probe graph and silently breaking all existing callers; (2) graph_debug.h is self-labeled temporary yet registers an atexit handler and atomic counters in every build.

transformer_engine/common/fused_attn/fused_attn.cpp (deprecated wrapper batch_size=0 regression), transformer_engine/common/fused_attn/graph_debug.h (unconditional atexit handler in production builds)

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn.cpp	Core backend selection rewritten to use probe approach; deprecated wrapper silently breaks for all callers because batch_size=0 causes cuDNN to reject every probe graph
transformer_engine/common/fused_attn/config_and_params.h	New FusedAttnConfig/FwdParams/BwdParams structs; defaults is_training=true, batch_size=0, o_format=NOT_SET expose silent misuse by the deprecated wrapper
transformer_engine/common/fused_attn/config_and_params.cpp	New 1156-line file implementing the opaque config/params handle API; getter/setter switch statements appear complete and reachable
transformer_engine/common/fused_attn/graph_debug.h	Self-described temporary debug instrumentation included unconditionally; atexit handler and atomic counters compiled into every build
transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu	Refactored to accept FusedAttnConfig; causal mask now uses set_diagonal_band_right_bound(0) via the new diagonal-alignment API; is_supported_f16_fwd/bwd probe functions added
transformer_engine/common/fused_attn/fused_attn_fp8.cu	Same FusedAttnConfig refactor as F16; new is_supported_fp8_fwd/bwd probe functions properly use std::to_string for error messages
transformer_engine/common/include/transformer_engine/fused_attn.h	Adds 749 lines of new public API (NVTEFusedAttnConfig, v2 entry-points); old nvte_get_fused_attn_backend signature preserved but now delegates to the probe-based path
transformer_engine/pytorch/csrc/extensions/attention.cpp	get_fused_attn_backend now takes py::object and populates FusedAttnConfigWrapper via chained setters; fwd/bwd rewritten to use FusedAttnFwdParamsWrapper/BwdParamsWrapper
transformer_engine/jax/cpp_extensions/attention.py	FusedAttnHelper gains batch_size, bottom_right_diagonal, attn_scale, and bias-shape fields; get_fused_attn_backend now returns (backend, message) tuple

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller as PyTorch/JAX/C++ Caller
    participant v2 as nvte_get_fused_attn_backend_v2
    participant pfwd as is_supported_f16_fwd
    participant pbwd as is_supported_f16_bwd
    participant ifwd as fused_attn_arbitrary_seqlen_fwd_impl
    participant ibwd as fused_attn_arbitrary_seqlen_bwd_impl
    participant cuDNN as cuDNN Graph Build

    Caller->>v2: NVTEFusedAttnConfig
    v2->>v2: early dtype/heads/seqlen checks
    v2->>pfwd: cfg + handle
    pfwd->>ifwd: "cfg with all device ptrs = nullptr"
    ifwd->>cuDNN: build fwd graph (batch, s_q, s_kv, h, d)
    alt graph builds OK
        cuDNN-->>ifwd: compiled graph
        ifwd-->>pfwd: workspace_size
        pfwd-->>v2: empty string (supported)
    else cuDNN rejects
        cuDNN-->>ifwd: exception
        pfwd-->>v2: error string
        v2-->>Caller: NVTE_No_Backend
    end
    alt is_training and not is_forward
        v2->>pbwd: cfg + handle
        pbwd->>ibwd: "cfg with all device ptrs = nullptr"
        ibwd->>cuDNN: build bwd graph
        cuDNN-->>ibwd: result or exception
        pbwd-->>v2: empty string or error string
    end
    v2-->>Caller: NVTE_F16_arbitrary_seqlen or NVTE_No_Backend

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Caller as PyTorch/JAX/C++ Caller
    participant v2 as nvte_get_fused_attn_backend_v2
    participant pfwd as is_supported_f16_fwd
    participant pbwd as is_supported_f16_bwd
    participant ifwd as fused_attn_arbitrary_seqlen_fwd_impl
    participant ibwd as fused_attn_arbitrary_seqlen_bwd_impl
    participant cuDNN as cuDNN Graph Build

    Caller->>v2: NVTEFusedAttnConfig
    v2->>v2: early dtype/heads/seqlen checks
    v2->>pfwd: cfg + handle
    pfwd->>ifwd: "cfg with all device ptrs = nullptr"
    ifwd->>cuDNN: build fwd graph (batch, s_q, s_kv, h, d)
    alt graph builds OK
        cuDNN-->>ifwd: compiled graph
        ifwd-->>pfwd: workspace_size
        pfwd-->>v2: empty string (supported)
    else cuDNN rejects
        cuDNN-->>ifwd: exception
        pfwd-->>v2: error string
        v2-->>Caller: NVTE_No_Backend
    end
    alt is_training and not is_forward
        v2->>pbwd: cfg + handle
        pbwd->>ibwd: "cfg with all device ptrs = nullptr"
        ibwd->>cuDNN: build bwd graph
        cuDNN-->>ibwd: result or exception
        pbwd-->>v2: empty string or error string
    end
    v2-->>Caller: NVTE_F16_arbitrary_seqlen or NVTE_No_Backend

_{Reviews (20): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

for more information, see https://pre-commit.ci

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

for more information, see https://pre-commit.ci

cyanguwa · 2026-05-08T01:35:11Z

/te-ci L1

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

for more information, see https://pre-commit.ci

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

for more information, see https://pre-commit.ci

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

pggPL · 2026-07-08T11:26:16Z

There is a limitation for torch.compile dynamic shapes here. If b and s are passed to black-box function affecting the program flow, I will need to force recompilation on every b or s change.
If we have args like s_greater_or_eq_512 and we check s >= 512 in python then no recompilation would be needed for every change of s. It may be more problematic for b.

I don't think this is a blocker for this PR, but we need to keep it in mind.

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

… fwd/bwd, thread bias shapes through in jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

for more information, see https://pre-commit.ci

…cache_key, fix Jax bias Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

…raph_debug, fix bias shape handling Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

greptile-apps · 2026-07-17T10:18:50Z

+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+// ============================================================================
+// [GRAPH-DEBUG] TEMPORARY DEBUG INSTRUMENTATION -- REMOVE AFTER VERIFICATION.
+//
+// Counts fused-attention cuDNN graph *builds* (cache misses that construct a new
+// graph) vs. *executions* (real forward/backward runs, excluding workspace-sizing
+// probes) to detect redundant graph construction.
+//
+// Enable at runtime with:  export NVTE_FUSED_ATTN_GRAPH_DEBUG=1
+// A running "BUILD" line is printed whenever a new graph is constructed, and a
+// "SUMMARY" line with final totals is printed at process exit.
+//
+// To remove all of this instrumentation later:
+//   1. Delete this file (graph_debug.h).
+//   2. Remove every line tagged with the "[GRAPH-DEBUG]" marker in:
+//        - fused_attn_fp8.cu
+//        - fused_attn_f16_arbitrary_seqlen.cu
+// ============================================================================


Temporary debug instrumentation shipped in production build

This file is explicitly labeled "TEMPORARY DEBUG INSTRUMENTATION -- REMOVE AFTER VERIFICATION" in its own header comment, yet it is included in this PR. The atexit handler and atomic counters are compiled into every build unconditionally; only the runtime output is gated behind NVTE_FUSED_ATTN_GRAPH_DEBUG. This should be removed before merge or tracked with a follow-up issue — otherwise the instrumentation accumulates permanently in the library.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

for more information, see https://pre-commit.ci

greptile-apps · 2026-07-17T16:53:59Z

+  transformer_engine::FusedAttnConfig cfg{};
+  cfg.qkv_layout = qkv_layout;
+  cfg.bias_type = bias_type;


Deprecated wrapper leaves batch_size=0, breaking all BSHD/SBHD probes

FusedAttnConfig cfg{} zero-initialises batch_size. The old API signature had no such parameter, so the wrapper never sets it. is_supported_f16_fwd then calls derive() (which bucketed_batch_size inherits from zero) and calls into fused_attn_arbitrary_seqlen_fwd_impl with b = 0. cuDNN rejects a graph whose batch tensor has size 0, throwing an exception. The catch block returns a non-empty error string, and nvte_get_fused_attn_backend_v2 returns NVTE_No_Backend for every BSHD/SBHD F16 configuration — a silent regression for every existing caller of the deprecated entry-point. A safe fallback is to seed cfg.batch_size = 1 since the probe only needs a structurally-valid (non-zero) dimension to succeed.

Suggested change

transformer_engine::FusedAttnConfig cfg{};

cfg.qkv_layout = qkv_layout;

cfg.bias_type = bias_type;

transformer_engine::FusedAttnConfig cfg{};

cfg.batch_size = 1; // probe needs a non-zero batch; old API had no batch_size parameter

cfg.qkv_layout = qkv_layout;

cfg.bias_type = bias_type;

cyanguwa and others added 4 commits May 5, 2026 18:55

refactor nvte_get_fused_attn_backend with FE calls

2b9fbc5

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

16b837c

for more information, see https://pre-commit.ci

Merge branch 'main' into fe_check_support

5a482f9

replace code+string with string only

42bcd89

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa mentioned this pull request May 7, 2026

[PyTorch/Common] Remove legacy FP8DS implementation #2959

Merged

13 tasks

clean up logic/comments/structure

de8e814

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa changed the title ~~[Common] Refactor nvte_get_fused_attn_backend with cudnn-frontend calls~~ [All] Refactor nvte_get_fused_attn_backend with cudnn-frontend calls May 8, 2026

Merge branch 'main' into fe_check_support

4b8c7ed

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa marked this pull request as ready for review May 8, 2026 00:10

pre-commit-ci Bot and others added 2 commits May 8, 2026 00:11

[pre-commit.ci] auto fixes from pre-commit.com hooks

81e59a9

for more information, see https://pre-commit.ci

Merge branch 'main' into fe_check_support

5640c68

greptile-apps Bot reviewed May 8, 2026

View reviewed changes

cyanguwa and others added 2 commits May 7, 2026 17:22

fix compilation errors

6c5126d

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

d35bff7

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed May 8, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp Outdated

cyanguwa and others added 3 commits May 7, 2026 18:30

remove handle from API; add bottom_right_diagonal

f6fc585

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

add batch_size to API

3e666b0

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

056aba6

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed May 8, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated

Comment thread transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp

fix jax binding

e054863

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

greptile-apps Bot reviewed May 8, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated

Comment thread transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated

cyanguwa and others added 3 commits May 7, 2026 22:28

specify o_dtype for FP8s

a7fe928

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix BRCM and custom_fp8 tests

c9b22b5

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

ac44e66

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed May 8, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp Outdated

Comment thread transformer_engine/jax/cpp_extensions/attention.py

cyanguwa and others added 2 commits May 8, 2026 12:19

add o_format/etc to API and other tweaks

9131b2d

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

956f159

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed May 8, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp Outdated

greptile-apps Bot reviewed May 12, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp Outdated

Merge branch 'main' into fe_check_support

3bed68d

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa requested review from KshitijLakhani, jberchtold-nvidia, ksivaman and ptrendx as code owners July 7, 2026 10:37

jberchtold-nvidia reviewed Jul 7, 2026

View reviewed changes

Comment thread transformer_engine/jax/flax/transformer.py

Comment thread transformer_engine/jax/csrc/extensions/pybind.cpp Outdated

Merge branch 'main' into fe_check_support

d634f3b

greptile-apps Bot reviewed Jul 8, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp Outdated

cyanguwa added 2 commits July 11, 2026 08:37

replace with opaque handle

1c09072

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

repeat with fwd/bwd params

ac19f9d

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

greptile-apps Bot reviewed Jul 14, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/config_and_params.cpp Outdated

remove bucketed b/t_q/t_kv

88a327c

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

greptile-apps Bot reviewed Jul 14, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/config_and_params.cpp Outdated

Comment thread transformer_engine/jax/cpp_extensions/attention.py

thread _v2 through, fix default scaling mode, make config specific to…

391fe2e

… fwd/bwd, thread bias shapes through in jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa force-pushed the fe_check_support branch from 57bc228 to 391fe2e Compare July 14, 2026 21:33

cyanguwa and others added 2 commits July 14, 2026 14:34

Merge branch 'main' into fe_check_support

877bf99

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

261bb9a

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed Jul 14, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/config_and_params.cpp Outdated

Comment thread transformer_engine/common/fused_attn/config_and_params.cpp Outdated

cyanguwa added 2 commits July 15, 2026 07:17

reorder struct fields, consolidate APIs to derive, make_config, make_…

c03d852

…cache_key, fix Jax bias Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

simplify fused attn config/params wrappers via set_attr helper, add g…

eadd005

…raph_debug, fix bias shape handling Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa force-pushed the fe_check_support branch from f5ba9d5 to eadd005 Compare July 16, 2026 18:18

greptile-apps Bot reviewed Jul 16, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp

Merge branch 'main' into fe_check_support

d052fcd

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

greptile-apps Bot reviewed Jul 17, 2026

View reviewed changes

clean up derived fields, debug probe/exec graph mismatches

28f5a8c

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa force-pushed the fe_check_support branch from c2152da to 28f5a8c Compare July 17, 2026 16:36

[pre-commit.ci] auto fixes from pre-commit.com hooks

4419cf8

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed Jul 17, 2026

View reviewed changes

Uh oh!

Conversation

cyanguwa commented May 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented May 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 2/5

Important Files Changed

Sequence Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cyanguwa commented May 8, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pggPL commented Jul 8, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

cyanguwa commented May 6, 2026 •

edited

Loading

greptile-apps Bot commented May 8, 2026 •

edited

Loading